{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "生成样本网格...\n",
      "生成数据增强示例...\n",
      "生成数字分布图...\n",
      "生成像素强度分布图...\n",
      "所有可视化完成！\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from torch.utils.data import DataLoader\n",
    "# 设置中文字体\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "plt.rcParams['axes.unicode_minus'] = False \n",
    "class MNISTVisualizer:\n",
    "    def __init__(self):\n",
    "        # 数据预处理\n",
    "        transform = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "        \n",
    "        # 加载数据集\n",
    "        self.train_dataset = torchvision.datasets.MNIST(\n",
    "            root='./data',\n",
    "            train=True,\n",
    "            download=True,\n",
    "            transform=transform\n",
    "        )\n",
    "        \n",
    "        self.test_dataset = torchvision.datasets.MNIST(\n",
    "            root='./data',\n",
    "            train=False,\n",
    "            download=True,\n",
    "            transform=transform\n",
    "        )\n",
    "        \n",
    "        # 创建数据加载器\n",
    "        self.train_loader = DataLoader(self.train_dataset, batch_size=64, shuffle=True)\n",
    "        self.test_loader = DataLoader(self.test_dataset, batch_size=64, shuffle=False)\n",
    "\n",
    "    def show_sample_grid(self, num_samples=25):\n",
    "        \"\"\"显示样本图片网格\"\"\"\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        for i in range(num_samples):\n",
    "            plt.subplot(5, 5, i + 1)\n",
    "            img, label = self.train_dataset[i]\n",
    "            plt.imshow(img.squeeze(), cmap='gray')\n",
    "            plt.title(f'标签: {label}')\n",
    "            plt.axis('off')\n",
    "        plt.tight_layout()\n",
    "        plt.savefig('sample_grid.png')\n",
    "        plt.close()\n",
    "\n",
    "    def show_augmented_samples(self):\n",
    "        \"\"\"显示数据增强效果\"\"\"\n",
    "        # 定义多种数据增强\n",
    "        augmentations = [\n",
    "            ('原始图片', transforms.ToTensor()),\n",
    "            ('旋转', transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                transforms.RandomRotation(30)\n",
    "            ])),\n",
    "            ('平移', transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                transforms.RandomAffine(0, translate=(0.2, 0.2))\n",
    "            ])),\n",
    "            ('缩放', transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                transforms.RandomAffine(0, scale=(0.8, 1.2))\n",
    "            ]))\n",
    "        ]\n",
    "        \n",
    "        plt.figure(figsize=(15, 3))\n",
    "        img, _ = self.train_dataset[0]\n",
    "        for idx, (name, transform) in enumerate(augmentations):\n",
    "            plt.subplot(1, 4, idx + 1)\n",
    "            if idx == 0:\n",
    "                plt.imshow(img.squeeze(), cmap='gray')\n",
    "            else:\n",
    "                aug_img = transform(img.squeeze().numpy())\n",
    "                plt.imshow(aug_img.squeeze(), cmap='gray')\n",
    "            plt.title(name)\n",
    "            plt.axis('off')\n",
    "        plt.tight_layout()\n",
    "        plt.savefig('augmented_samples.png')\n",
    "        plt.close()\n",
    "\n",
    "    def plot_digit_distribution(self):\n",
    "        \"\"\"绘制数字分布直方图\"\"\"\n",
    "        labels = [label for _, label in self.train_dataset]\n",
    "        plt.figure(figsize=(10, 5))\n",
    "        plt.hist(labels, bins=10, rwidth=0.8, align='left')\n",
    "        plt.title('MNIST数据集数字分布')\n",
    "        plt.xlabel('数字')\n",
    "        plt.ylabel('数量')\n",
    "        plt.xticks(range(10))\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        plt.savefig('digit_distribution.png')\n",
    "        plt.close()\n",
    "\n",
    "    def plot_intensity_distribution(self):\n",
    "        \"\"\"绘制像素强度分布\"\"\"\n",
    "        intensities = []\n",
    "        for i in range(1000):  # 使用前1000张图片\n",
    "            img, _ = self.train_dataset[i]\n",
    "            intensities.extend(img.numpy().flatten())\n",
    "            \n",
    "        plt.figure(figsize=(10, 5))\n",
    "        plt.hist(intensities, bins=50, density=True)\n",
    "        plt.title('像素强度分布')\n",
    "        plt.xlabel('像素值')\n",
    "        plt.ylabel('密度')\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        plt.savefig('intensity_distribution.png')\n",
    "        plt.close()\n",
    "\n",
    "    def plot_confusion_matrix(self, model, device):\n",
    "        \"\"\"绘制混淆矩阵\"\"\"\n",
    "        model.eval()\n",
    "        all_preds = []\n",
    "        all_labels = []\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for data, labels in self.test_loader:\n",
    "                data = data.to(device)\n",
    "                outputs = model(data)\n",
    "                _, predicted = torch.max(outputs, 1)\n",
    "                all_preds.extend(predicted.cpu().numpy())\n",
    "                all_labels.extend(labels.numpy())\n",
    "        \n",
    "        # 计算混淆矩阵\n",
    "        cm = confusion_matrix(all_labels, all_preds)\n",
    "        \n",
    "        # 绘制热力图\n",
    "        plt.figure(figsize=(10, 8))\n",
    "        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')\n",
    "        plt.title('预测结果混淆矩阵')\n",
    "        plt.xlabel('预测标签')\n",
    "        plt.ylabel('真实标签')\n",
    "        plt.savefig('confusion_matrix.png')\n",
    "        plt.close()\n",
    "\n",
    "    def visualize_feature_maps(self, model, device):\n",
    "        \"\"\"可视化特征图\"\"\"\n",
    "        # 获取一个批次的数据\n",
    "        data, _ = next(iter(self.test_loader))\n",
    "        img = data[0:1].to(device)  # 只使用第一张图片\n",
    "        \n",
    "        # 获取第一个卷积层的输出\n",
    "        activation = {}\n",
    "        def get_activation(name):\n",
    "            def hook(model, input, output):\n",
    "                activation[name] = output.detach()\n",
    "            return hook\n",
    "        \n",
    "        # 注册钩子\n",
    "        model.features[0].register_forward_hook(get_activation('conv1'))\n",
    "        \n",
    "        # 前向传播\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            output = model(img)\n",
    "        \n",
    "        # 可视化特征图\n",
    "        plt.figure(figsize=(15, 5))\n",
    "        \n",
    "        # 原始图片\n",
    "        plt.subplot(1, 3, 1)\n",
    "        plt.imshow(img.cpu().squeeze(), cmap='gray')\n",
    "        plt.title('原始图片')\n",
    "        plt.axis('off')\n",
    "        \n",
    "        # 特征图\n",
    "        feature_maps = activation['conv1'].squeeze().cpu()\n",
    "        plt.subplot(1, 3, 2)\n",
    "        plt.imshow(feature_maps[0], cmap='viridis')\n",
    "        plt.title('特征图 1')\n",
    "        plt.axis('off')\n",
    "        \n",
    "        plt.subplot(1, 3, 3)\n",
    "        plt.imshow(feature_maps[1], cmap='viridis')\n",
    "        plt.title('特征图 2')\n",
    "        plt.axis('off')\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        plt.savefig('feature_maps.png')\n",
    "        plt.close()\n",
    "\n",
    "    def plot_learning_curves(self, train_losses, test_losses, train_accs, test_accs):\n",
    "        \"\"\"绘制学习曲线\"\"\"\n",
    "        plt.figure(figsize=(15, 5))\n",
    "        \n",
    "        # 损失曲线\n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.plot(train_losses, label='训练损失')\n",
    "        plt.plot(test_losses, label='测试损失')\n",
    "        plt.title('损失曲线')\n",
    "        plt.xlabel('轮次')\n",
    "        plt.ylabel('损失')\n",
    "        plt.legend()\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        \n",
    "        # 准确率曲线\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.plot(train_accs, label='训练准确率')\n",
    "        plt.plot(test_accs, label='测试准确率')\n",
    "        plt.title('准确率曲线')\n",
    "        plt.xlabel('轮次')\n",
    "        plt.ylabel('准确率 (%)')\n",
    "        plt.legend()\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        plt.savefig('learning_curves.png')\n",
    "        plt.close()\n",
    "\n",
    "def main():\n",
    "    # 创建可视化器\n",
    "    visualizer = MNISTVisualizer()\n",
    "    \n",
    "    print(\"生成样本网格...\")\n",
    "    visualizer.show_sample_grid()\n",
    "    \n",
    "    print(\"生成数据增强示例...\")\n",
    "    visualizer.show_augmented_samples()\n",
    "    \n",
    "    print(\"生成数字分布图...\")\n",
    "    visualizer.plot_digit_distribution()\n",
    "    \n",
    "    print(\"生成像素强度分布图...\")\n",
    "    visualizer.plot_intensity_distribution()\n",
    "    \n",
    "    \n",
    "    \n",
    "    print(\"所有可视化完成！\")\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "使用设备: cuda\n",
      "\n",
      "开始训练...\n",
      "\n",
      "第 1 轮训练:\n",
      "Loss: 0.576 | Acc: 81.80%\n",
      "Loss: 0.178 | Acc: 88.16%\n",
      "Loss: 0.127 | Acc: 90.69%\n",
      "Loss: 0.100 | Acc: 92.23%\n",
      "Loss: 0.096 | Acc: 93.21%\n",
      "Loss: 0.080 | Acc: 93.91%\n",
      "Loss: 0.073 | Acc: 94.45%\n",
      "Loss: 0.071 | Acc: 94.88%\n",
      "Loss: 0.072 | Acc: 95.21%\n",
      "训练损失: 0.003 | 训练准确率: 95.32%\n",
      "测试损失: 0.037 | 测试准确率: 98.87%\n",
      "保存最佳模型...\n",
      "\n",
      "第 2 轮训练:\n",
      "Loss: 0.071 | Acc: 97.95%\n",
      "Loss: 0.059 | Acc: 98.10%\n",
      "Loss: 0.058 | Acc: 98.10%\n",
      "Loss: 0.046 | Acc: 98.24%\n",
      "Loss: 0.053 | Acc: 98.28%\n",
      "Loss: 0.044 | Acc: 98.33%\n",
      "Loss: 0.049 | Acc: 98.35%\n",
      "Loss: 0.047 | Acc: 98.39%\n",
      "Loss: 0.057 | Acc: 98.39%\n",
      "训练损失: 0.002 | 训练准确率: 98.39%\n",
      "测试损失: 0.029 | 测试准确率: 99.05%\n",
      "保存最佳模型...\n",
      "\n",
      "第 3 轮训练:\n",
      "Loss: 0.028 | Acc: 98.98%\n",
      "Loss: 0.043 | Acc: 98.84%\n",
      "Loss: 0.039 | Acc: 98.81%\n",
      "Loss: 0.034 | Acc: 98.83%\n",
      "Loss: 0.039 | Acc: 98.81%\n",
      "Loss: 0.049 | Acc: 98.77%\n",
      "Loss: 0.043 | Acc: 98.74%\n",
      "Loss: 0.034 | Acc: 98.75%\n",
      "Loss: 0.040 | Acc: 98.77%\n",
      "训练损失: 0.001 | 训练准确率: 98.77%\n",
      "测试损失: 0.025 | 测试准确率: 99.21%\n",
      "保存最佳模型...\n",
      "\n",
      "第 4 轮训练:\n",
      "Loss: 0.029 | Acc: 99.02%\n",
      "Loss: 0.028 | Acc: 99.06%\n",
      "Loss: 0.029 | Acc: 99.04%\n",
      "Loss: 0.026 | Acc: 99.07%\n",
      "Loss: 0.028 | Acc: 99.07%\n",
      "Loss: 0.035 | Acc: 99.05%\n",
      "Loss: 0.028 | Acc: 99.06%\n",
      "Loss: 0.026 | Acc: 99.06%\n",
      "Loss: 0.035 | Acc: 99.04%\n",
      "训练损失: 0.001 | 训练准确率: 99.04%\n",
      "测试损失: 0.023 | 测试准确率: 99.21%\n",
      "\n",
      "第 5 轮训练:\n",
      "Loss: 0.019 | Acc: 99.44%\n",
      "Loss: 0.024 | Acc: 99.34%\n",
      "Loss: 0.024 | Acc: 99.35%\n",
      "Loss: 0.026 | Acc: 99.35%\n",
      "Loss: 0.025 | Acc: 99.33%\n",
      "Loss: 0.023 | Acc: 99.31%\n",
      "Loss: 0.023 | Acc: 99.29%\n",
      "Loss: 0.031 | Acc: 99.27%\n",
      "Loss: 0.031 | Acc: 99.26%\n",
      "训练损失: 0.001 | 训练准确率: 99.25%\n",
      "测试损失: 0.026 | 测试准确率: 99.24%\n",
      "保存最佳模型...\n",
      "\n",
      "第 6 轮训练:\n",
      "Loss: 0.024 | Acc: 99.41%\n",
      "Loss: 0.022 | Acc: 99.38%\n",
      "Loss: 0.015 | Acc: 99.44%\n",
      "Loss: 0.015 | Acc: 99.45%\n",
      "Loss: 0.024 | Acc: 99.40%\n",
      "Loss: 0.020 | Acc: 99.40%\n",
      "Loss: 0.017 | Acc: 99.40%\n",
      "Loss: 0.018 | Acc: 99.40%\n",
      "Loss: 0.023 | Acc: 99.39%\n",
      "训练损失: 0.001 | 训练准确率: 99.38%\n",
      "测试损失: 0.025 | 测试准确率: 99.23%\n",
      "\n",
      "第 7 轮训练:\n",
      "Loss: 0.015 | Acc: 99.47%\n",
      "Loss: 0.018 | Acc: 99.43%\n",
      "Loss: 0.017 | Acc: 99.45%\n",
      "Loss: 0.018 | Acc: 99.45%\n",
      "Loss: 0.017 | Acc: 99.46%\n",
      "Loss: 0.015 | Acc: 99.48%\n",
      "Loss: 0.019 | Acc: 99.48%\n",
      "Loss: 0.017 | Acc: 99.47%\n",
      "Loss: 0.014 | Acc: 99.48%\n",
      "训练损失: 0.001 | 训练准确率: 99.47%\n",
      "测试损失: 0.023 | 测试准确率: 99.30%\n",
      "保存最佳模型...\n",
      "\n",
      "第 8 轮训练:\n",
      "Loss: 0.013 | Acc: 99.52%\n",
      "Loss: 0.010 | Acc: 99.62%\n",
      "Loss: 0.017 | Acc: 99.55%\n",
      "Loss: 0.011 | Acc: 99.58%\n",
      "Loss: 0.011 | Acc: 99.58%\n",
      "Loss: 0.022 | Acc: 99.56%\n",
      "Loss: 0.016 | Acc: 99.54%\n",
      "Loss: 0.020 | Acc: 99.51%\n",
      "Loss: 0.015 | Acc: 99.52%\n",
      "训练损失: 0.001 | 训练准确率: 99.52%\n",
      "测试损失: 0.030 | 测试准确率: 99.28%\n",
      "\n",
      "第 9 轮训练:\n",
      "Loss: 0.013 | Acc: 99.66%\n",
      "Loss: 0.015 | Acc: 99.56%\n",
      "Loss: 0.016 | Acc: 99.54%\n",
      "Loss: 0.007 | Acc: 99.59%\n",
      "Loss: 0.013 | Acc: 99.57%\n",
      "Loss: 0.012 | Acc: 99.58%\n",
      "Loss: 0.013 | Acc: 99.58%\n",
      "Loss: 0.017 | Acc: 99.56%\n",
      "Loss: 0.018 | Acc: 99.55%\n",
      "训练损失: 0.001 | 训练准确率: 99.54%\n",
      "测试损失: 0.028 | 测试准确率: 99.23%\n",
      "\n",
      "第 10 轮训练:\n",
      "Loss: 0.014 | Acc: 99.55%\n",
      "Loss: 0.010 | Acc: 99.62%\n",
      "Loss: 0.008 | Acc: 99.67%\n",
      "Loss: 0.009 | Acc: 99.68%\n",
      "Loss: 0.010 | Acc: 99.67%\n",
      "Loss: 0.011 | Acc: 99.66%\n",
      "Loss: 0.009 | Acc: 99.66%\n",
      "Loss: 0.017 | Acc: 99.63%\n",
      "Loss: 0.014 | Acc: 99.62%\n",
      "训练损失: 0.001 | 训练准确率: 99.62%\n",
      "测试损失: 0.025 | 测试准确率: 99.27%\n",
      "\n",
      "训练完成！\n",
      "最佳测试准确率: 99.30%\n",
      "\n",
      "生成可视化结果...\n",
      "所有可视化已保存！\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# 定义神经网络\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "            nn.Conv2d(32, 64, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(64 * 7 * 7, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(512, 10)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(-1, 64 * 7 * 7)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "class MNISTTrainer:\n",
    "    def __init__(self, device):\n",
    "        self.device = device\n",
    "        # 数据预处理\n",
    "        transform = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "        \n",
    "        # 加载数据集\n",
    "        self.train_dataset = torchvision.datasets.MNIST(\n",
    "            root='./data', train=True, download=True, transform=transform)\n",
    "        self.test_dataset = torchvision.datasets.MNIST(\n",
    "            root='./data', train=False, transform=transform)\n",
    "        \n",
    "        # 创建数据加载器\n",
    "        self.train_loader = DataLoader(self.train_dataset, batch_size=64, shuffle=True)\n",
    "        self.test_loader = DataLoader(self.test_dataset, batch_size=64, shuffle=False)\n",
    "        \n",
    "        # 创建模型\n",
    "        self.model = Net().to(device)\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "        self.optimizer = optim.Adam(self.model.parameters())\n",
    "        \n",
    "        # 记录训练过程\n",
    "        self.train_losses = []\n",
    "        self.test_losses = []\n",
    "        self.train_accs = []\n",
    "        self.test_accs = []\n",
    "\n",
    "    def train_epoch(self):\n",
    "        self.model.train()\n",
    "        total_loss = 0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        \n",
    "        for batch_idx, (data, target) in enumerate(self.train_loader):\n",
    "            data, target = data.to(self.device), target.to(self.device)\n",
    "            \n",
    "            self.optimizer.zero_grad()\n",
    "            output = self.model(data)\n",
    "            loss = self.criterion(output, target)\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            _, predicted = output.max(1)\n",
    "            total += target.size(0)\n",
    "            correct += predicted.eq(target).sum().item()\n",
    "            \n",
    "            if batch_idx % 100 == 99:\n",
    "                print(f'Loss: {total_loss/100:.3f} | Acc: {100.*correct/total:.2f}%')\n",
    "                total_loss = 0\n",
    "                \n",
    "        return total_loss / len(self.train_loader), 100. * correct / total\n",
    "\n",
    "    def test(self):\n",
    "        self.model.eval()\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for data, target in self.test_loader:\n",
    "                data, target = data.to(self.device), target.to(self.device)\n",
    "                output = self.model(data)\n",
    "                test_loss += self.criterion(output, target).item()\n",
    "                _, predicted = output.max(1)\n",
    "                total += target.size(0)\n",
    "                correct += predicted.eq(target).sum().item()\n",
    "                \n",
    "        test_loss /= len(self.test_loader)\n",
    "        accuracy = 100. * correct / total\n",
    "        \n",
    "        return test_loss, accuracy\n",
    "\n",
    "    def visualize_results(self):\n",
    "        # 1. 学习曲线\n",
    "        plt.figure(figsize=(15, 5))\n",
    "        \n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.plot(self.train_losses, label='训练损失')\n",
    "        plt.plot(self.test_losses, label='测试损失')\n",
    "        plt.title('损失曲线')\n",
    "        plt.xlabel('轮次')\n",
    "        plt.ylabel('损失')\n",
    "        plt.legend()\n",
    "        plt.grid(True)\n",
    "        \n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.plot(self.train_accs, label='训练准确率')\n",
    "        plt.plot(self.test_accs, label='测试准确率')\n",
    "        plt.title('准确率曲线')\n",
    "        plt.xlabel('轮次')\n",
    "        plt.ylabel('准确率 (%)')\n",
    "        plt.legend()\n",
    "        plt.grid(True)\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        plt.savefig('learning_curves.png')\n",
    "        plt.close()\n",
    "        \n",
    "        # 2. 混淆矩阵\n",
    "        all_preds = []\n",
    "        all_labels = []\n",
    "        self.model.eval()\n",
    "        with torch.no_grad():\n",
    "            for data, labels in self.test_loader:\n",
    "                data = data.to(self.device)\n",
    "                outputs = self.model(data)\n",
    "                _, predicted = torch.max(outputs, 1)\n",
    "                all_preds.extend(predicted.cpu().numpy())\n",
    "                all_labels.extend(labels.numpy())\n",
    "        \n",
    "        cm = confusion_matrix(all_labels, all_preds)\n",
    "        plt.figure(figsize=(10, 8))\n",
    "        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')\n",
    "        plt.title('预测结果混淆矩阵')\n",
    "        plt.xlabel('预测标签')\n",
    "        plt.ylabel('真实标签')\n",
    "        plt.savefig('confusion_matrix.png')\n",
    "        plt.close()\n",
    "        \n",
    "        # 3. 预测结果示例\n",
    "        plt.figure(figsize=(15, 5))\n",
    "        with torch.no_grad():\n",
    "            data, labels = next(iter(self.test_loader))\n",
    "            data = data.to(self.device)\n",
    "            outputs = self.model(data)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            \n",
    "            for i in range(10):\n",
    "                plt.subplot(2, 5, i + 1)\n",
    "                plt.imshow(data[i].cpu().squeeze(), cmap='gray')\n",
    "                plt.title(f'预测: {predicted[i].item()}\\n实际: {labels[i].item()}')\n",
    "                plt.axis('off')\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        plt.savefig('prediction_examples.png')\n",
    "        plt.close()\n",
    "\n",
    "def main():\n",
    "    # 设置设备\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(f\"使用设备: {device}\")\n",
    "    \n",
    "    # 创建训练器\n",
    "    trainer = MNISTTrainer(device)\n",
    "    \n",
    "    # 训练模型\n",
    "    num_epochs = 10\n",
    "    best_acc = 0\n",
    "    \n",
    "    print(\"\\n开始训练...\")\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f'\\n第 {epoch+1} 轮训练:')\n",
    "        \n",
    "        # 训练一轮\n",
    "        train_loss, train_acc = trainer.train_epoch()\n",
    "        trainer.train_losses.append(train_loss)\n",
    "        trainer.train_accs.append(train_acc)\n",
    "        \n",
    "        # 测试\n",
    "        test_loss, test_acc = trainer.test()\n",
    "        trainer.test_losses.append(test_loss)\n",
    "        trainer.test_accs.append(test_acc)\n",
    "        \n",
    "        print(f'训练损失: {train_loss:.3f} | 训练准确率: {train_acc:.2f}%')\n",
    "        print(f'测试损失: {test_loss:.3f} | 测试准确率: {test_acc:.2f}%')\n",
    "        \n",
    "        # 保存最佳模型\n",
    "        if test_acc > best_acc:\n",
    "            print('保存最佳模型...')\n",
    "            torch.save(trainer.model.state_dict(), 'best_model.pth')\n",
    "            best_acc = test_acc\n",
    "    \n",
    "    print(\"\\n训练完成！\")\n",
    "    print(f\"最佳测试准确率: {best_acc:.2f}%\")\n",
    "    \n",
    "    # 可视化结果\n",
    "    print(\"\\n生成可视化结果...\")\n",
    "    trainer.visualize_results()\n",
    "    print(\"所有可视化已保存！\")\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
